package org.longteng.sensitive.mybatis.intercepts;

import cn.hutool.core.util.ReflectUtil;
import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import org.longteng.sensitive.annotation.EncryptField;
import org.longteng.sensitive.ISensitiveDataEncrypt;
import org.longteng.sensitive.utils.SensitiveUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.type.UnknownTypeHandler;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.lang.annotation.Annotation;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.sql.PreparedStatement;
import java.util.*;


/**
 * 参数加密
 * 依赖注解 SensitiveData 与 SensitiveField
 * 需要加解密的类：类上需要标记SensitiveData 敏感属性 包括集合 需要添加SensitiveField
 * 解释: @Intercepts 注解开启拦截器，@Signature 注解定义拦截器的实际类型。
 *
 * @Signature中 type 属性指定当前拦截器使用StatementHandler 、ResultSetHandler、ParameterHandler，Executor的一种
 * method 属性指定使用以上四种类型的具体方法（可进入class内部查看其方法）。
 * args 属性指定预编译语句
 *
 * Interceptor 是Mybatis 的 拦截器接口
 */
@Intercepts({
        @Signature(type = ParameterHandler.class, method = "setParameters", args = {PreparedStatement.class})
})
@Component
@Slf4j
public class EncryptInterceptor implements Interceptor {
    /**注入数据加解密实现
     * <br>需要在启动应用中用@Bean构建ISensitiveDataEncrypt的实现类
     * */
    private ISensitiveDataEncrypt sensitiveDataEncrypt;
    @Resource
    public void setSensitiveDataEncrypt(ISensitiveDataEncrypt sensitiveDataEncrypt) {
        this.sensitiveDataEncrypt = sensitiveDataEncrypt;
        log.info("注入数据加密实现:{}", sensitiveDataEncrypt.getClass());
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        ParameterHandler parameterHandler = (ParameterHandler) invocation.getTarget();
        // MappedStatement ms = (MappedStatement) invocation.getArgs()[0];
        this.encryptSensitiveField(parameterHandler);

        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        /*
         * 这里就是将拦截器和对象包装在一起:
         * ①获取拦截器的Intercepts注解的所有Signature参数，即该拦截器要拦截的类和对象的方法、参数
         * ②获取拦截对象的类
         * ③获取所有接口
         * ④根据返回的接口数量，判断是否要拦截的，要拦截的对象生成将拦截器和拦截对象封装在一起的代理对象
         */
        return Plugin.wrap(target, this);
    }

    /**
     * 加密字段进行参数加密
     * @param parameterHandler
     */
    private void encryptSensitiveField(ParameterHandler parameterHandler) throws NoSuchFieldException, ClassNotFoundException {
        Object parameterObject = parameterHandler.getParameterObject();
        if (Objects.isNull(parameterObject)) {
            return;
        }
        if (parameterObject instanceof MapperMethod.ParamMap) {
            handleMethodInMapper(parameterHandler);
            return;
        }

        handleMethodNotInmapper(parameterObject);
    }

    /**
     * 参数集处理
     * @param parameterHandler
     * @throws NoSuchFieldException
     * @throws ClassNotFoundException
     */
    private void handleMethodInMapper(ParameterHandler parameterHandler) throws NoSuchFieldException, ClassNotFoundException {
        MapperMethod.ParamMap<Object> paramMap = ( MapperMethod.ParamMap<Object>) parameterHandler.getParameterObject();
        for (Map.Entry<String, Object> entry : paramMap.entrySet()) {
            String key = entry.getKey();
            Object paramValue = entry.getValue();
            if (!key.startsWith("param") && Objects.nonNull(paramValue)) {
                encrypt(parameterHandler, entry);
            }
        }
    }

    /**
     * 对DB字段进行加密
     *
     * @param parameterHandler
     * @param entry
     */
    private void encrypt(ParameterHandler parameterHandler, Map.Entry<String, Object> entry) throws NoSuchFieldException, ClassNotFoundException {
        Object paramValue = entry.getValue();
        //字符串
        if (paramValue instanceof String) {
            handleString(parameterHandler, entry);
            return;
        }
        //请求参数不是字符串
        log.debug("encrypt,请求参数不是字符串的处理.{}", paramValue==null?null:paramValue.getClass());
        handleMethodNotInmapper(paramValue);
    }

    /**
     * 参数对象处理 只处理：1.包含SensitiveData注解,属性包含EncryptField注解 2.属于com.atpingan的包或子包
     *
     * @param targetObj
     */
    private void handleMethodNotInmapper(Object targetObj) {
        if (Objects.isNull(targetObj)) {
            return;
        }
        Class<?> parameterClass = targetObj.getClass();
        Object entity = targetObj;
        Pair<Class<?>, Object> pair = handleQueryWrapper(targetObj);

        boolean isQueryWrapperOrLambdaQueryWrapper = Objects.nonNull(pair);
        if (isQueryWrapperOrLambdaQueryWrapper) {
            parameterClass = pair.getLeft();
            entity = pair.getRight();
        }
        if (Objects.isNull(entity)) {
            return;
        }
        //是否敏感类
        boolean sensitiveClass = SensitiveUtil.isSensitiveClass(parameterClass);
        if (!sensitiveClass) {
            return;
        }
        //路径前缀判断
        /*String typeName = parameterClass.getTypeName();
        if (!typeName.startsWith("com.atping")) {
            return;
        }
         */

        //获取敏感字段
        List<Field> sensitiveFields = SensitiveUtil.getSensitiveField(parameterClass);
        if (ObjectUtils.isEmpty(sensitiveFields)) {
            return;
        }
        encryptField(entity, sensitiveFields);
        if (isQueryWrapperOrLambdaQueryWrapper) {
            ReflectUtil.setFieldValue(targetObj, "entity", entity);
        }
    }

    /**
     * 参数字段 加密
     * @param entity
     * @param sensitiveFields
     */
    private void encryptField(Object entity, List<Field> sensitiveFields) {
        for (Field field : sensitiveFields) {  // 处理 有EncryptField.class的属性
            Object fieldValue = ReflectUtil.getFieldValue(entity, field);
            if (fieldValue instanceof String) {
                String fieldValueStr = (String) fieldValue;
                if(log.isDebugEnabled()) {
                    log.debug("字段{}参数在进行加密：{}", field.getName(),fieldValue);
                }
                TableField tableField = field.getAnnotation(TableField.class);
                if (Objects.nonNull(tableField)) {
                    if(! tableField.typeHandler().isAssignableFrom(UnknownTypeHandler.class)) {
                        if(log.isDebugEnabled()) {
                            log.debug("字段{}指定了字段Handler不在此处理加密：{}", field.getName(), tableField.typeHandler());
                        }
                        continue;
                    }
                }

                String fieldValueENCStr = (String) sensitiveDataEncrypt.encrypt(fieldValueStr, field);
                ReflectUtil.setFieldValue(entity, field, fieldValueENCStr);
            }
        }
    }

    /**
     * Lambda处理的参数
     * @param targetObj
     * @return
     */
    private Pair<Class<?>, Object> handleQueryWrapper(Object targetObj) {
        if (targetObj instanceof LambdaQueryWrapper) {
            return Pair.of(((LambdaQueryWrapper<?>) targetObj).getEntityClass(), ((LambdaQueryWrapper<?>) targetObj).getEntity());
        }
        if (targetObj instanceof LambdaUpdateWrapper) {
            return Pair.of(((LambdaUpdateWrapper<?>) targetObj).getEntityClass(), ((LambdaUpdateWrapper<?>) targetObj).getEntity());
        }
        if (targetObj instanceof QueryWrapper) {
            return Pair.of(((QueryWrapper<?>) targetObj).getEntityClass(), ((QueryWrapper<?>) targetObj).getEntity());
        }
        return null;
    }

    /**
     * string字段加密处理
     * @param parameterHandler
     * @param entry
     * @throws NoSuchFieldException
     * @throws ClassNotFoundException
     */
    private void handleString(ParameterHandler parameterHandler, Map.Entry<String, Object> entry) throws NoSuchFieldException, ClassNotFoundException {
        String key = entry.getKey();
        if(log.isDebugEnabled()) {
            log.debug("handleString(),sql字段参数名key:{},", key);
        }

        Object paramValue = entry.getValue();
        Class<? extends ParameterHandler> aClass = parameterHandler.getClass();
        Field mappedStatement = aClass.getDeclaredField("mappedStatement");
        ReflectUtil.setAccessible(mappedStatement);
        MappedStatement statement = (MappedStatement) ReflectUtil.getFieldValue(parameterHandler, mappedStatement);
        //方法命名空间
        String nameSpace = statement.getId();
        if (StringUtils.isBlank(nameSpace)) {
            return;
        }
        String methodName = nameSpace.substring(nameSpace.lastIndexOf(".") + 1);
        String className = nameSpace.substring(0, nameSpace.lastIndexOf("."));
        if(log.isDebugEnabled()) {
            log.debug("handleString(),取类{}的方法({})参数里面的请求参数注解列表.", className, methodName);
        }

        Method[] ms = Class.forName(className).getMethods();
        Optional<Method> optionalMethod = Arrays.stream(ms).filter(item -> StringUtils.equals(item.getName(), methodName)).findFirst();
        if (!optionalMethod.isPresent()) {
            return;
        }
        Method method = optionalMethod.get();
        ReflectUtil.setAccessible(method);
        //方法参数里面的请求参数注解列表
        Annotation[][] parameterAnnotations = method.getParameterAnnotations();
        boolean sensitiveField = Arrays.stream(parameterAnnotations).anyMatch(item -> Boolean.TRUE.equals(Arrays.stream(item).anyMatch(ite -> {
            if (ite instanceof EncryptField) {
                EncryptField sensitive = (EncryptField) ite;
                return StringUtils.equals(key, sensitive.value());    // 字段注解时指定字段名
            }
            return false;
        })));
        if (!sensitiveField) {
            return;
        }
        String encrypt = (String) sensitiveDataEncrypt.encrypt(paramValue);
        entry.setValue(encrypt);
    }

}
